#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@Project : aimix
@File    : batch_synthesis_or_ident_worker.py
@IDE     : PyCharm
@Author  : admin
@Date    : 2025/4/27 15:06
"""
import os
import json
import traceback
import tempfile
from PyQt5.QtCore import QThread, pyqtSignal
from utils.tools import request
from config import Config


class BatchSynthesisOrIdentWorker(QThread):
    """批量合成或识别音频线程"""
    result = pyqtSignal(str)
    failed = pyqtSignal(Exception)

    def __init__(self, content, headers, parent=None):
        super(BatchSynthesisOrIdentWorker, self).__init__()
        self.headers = headers
        self.content = json.loads(content)
        self.output_dir = self.content.get("output_dir", os.path.join(os.getcwd(), "output"))

    def format_timestamp(self, milliseconds: float) -> str:
        """格式化时间戳，包括毫秒"""
        total_seconds = milliseconds / 1000.0  # 将毫秒转换为秒
        m, s = divmod(total_seconds, 60)
        h, m = divmod(m, 60)
        s, ms = divmod(s, 1)  # 分离秒和毫秒
        return '{:02d}:{:02d}:{:02d},{:03d}'.format(int(h), int(m), int(s), int(ms * 1000))


    def generate_files_data(self, files):
        """
        生成文件上传所需的 files 参数
        :param files: 文件路径列表
        :return: files 参数
        """
        retv = []
        file_type_mapping = {
            '.mp4': 'video/mp4',
            '.mov': 'video/quicktime',
            '.avi': 'video/x-msvideo',
            '.mkv': 'video/x-matroska',
            '.mp3': 'audio/mpeg',
            '.wav': 'audio/wav',
            '.m4a': 'audio/m4a',
            '.aac': 'audio/aac'
        }

        for file_path in files:
            if os.path.exists(file_path):
                file_name = os.path.basename(file_path)
                file_extension = os.path.splitext(file_path)[1].lower()
                file_type = file_type_mapping.get(file_extension, 'application/octet-stream')
                retv.append(('files', (file_name, open(file_path, 'rb'), file_type)))
            else:
                print(f"文件不存在: {file_path}")
        return retv


    def identify(self, tasks):
        """识别音频"""
        files = self.generate_files_data(tasks)
        url = f"{Config.Host}/api/ident/speech/"
        self.headers.update({'Content-Type': 'application/json'})
        response = request(url, method='POST', headers=self.headers, data={}, files=files)
        response.raise_for_status()
        context = response.json()["transcripts"]
        for item in context:
            file_path = tempfile.mktemp(suffix='.srt', dir=self.output_dir)
            with open(file_path, 'w', encoding='utf-8') as f:
                for data in item:
                    f.write(str(data['sentence_id']) + '\n')
                    f.write(self.format_timestamp(data['start']) + '-->' + self.format_timestamp(data['end']) + '\n')
                    f.write(data['text'] + '\n\n')



    def synthesis(self, task_file, model, voice):
        """合成音频"""
        with open(task_file) as f:
            text = f.read()
        payload = json.dumps({
                "text": text,
                "model": model,
                "voice": voice,
                "suffix": "mp3",
                "volume": 100,
                "rate": 1.0,
                "pitch": 1.0

        })
        url = f"{Config.Host}/api/synthesis/audio/"
        self.headers.update({'Content-Type': 'application/json'})
        response = request(url, method='POST', headers=self.headers, data=payload)
        response.raise_for_status()
        file_url = response.json()["file_url"]
        audio_content = request(file_url).content
        temp_file = tempfile.mktemp(suffix=".mp3", dir=self.content["output_dir"])
        with open(temp_file, "wb") as f:
            f.write(audio_content)



    def run(self):
        # 开始批量合成或识别
        try:
            channel = self.content.get("channel")
            model, voice = self.content.get("voice").split(";")
            for task in self.content["files"]:
                task_file = task["file"]
                if channel == "audio" and task_file.endswith(".txt"):
                    self.synthesis(task_file, model, voice)
                elif channel == "subtitle" and task_file.endswith(".srt"):
                    self.identify([task_file])
                else:
                    print(f"不支持的文件类型: {channel} {task_file}")
                    continue
                self.result.emit(str(task["id"]))
        except Exception as e:
            self.failed.emit(e)
            traceback.print_exc()
